import numpy as np
from constants import DATA_ROOT


class DataConfig:
    def __init__(
            self,
            name,
            img_size,
            patch_size,
            num_channels=1,
            data_path='',
    ):
        super().__init__()
        self.name = name
        self.img_size = img_size
        self.patch_size = patch_size
        assert len(img_size) == len(patch_size)
        self.n_dim = len(img_size)
        self.grid_size = [s // p for s, p in zip(img_size, patch_size)]
        self.num_channels = num_channels
        self.token_dim = num_channels * np.prod(patch_size)
        self.seq_len = np.prod(self.grid_size)
        self.data_path = data_path

        self.modality_idx = None
        self.patch_proj_idx = None
        self.lo, self.hi = None, None

    def get_range(self):
        return self.lo, self.hi


PATCH_CONFIG = [
    (1, 16, 16, 16),  # 0
    (1, 4, 16, 16),  # 1
    (1, 16, 16),  # 2
    (3, 16, 16)  # 3
]
PATCH_RANGES = []
lo = 0
for p in PATCH_CONFIG:
    hi = lo + np.prod(p)
    PATCH_RANGES.append((lo, hi))
    lo = hi

MODALITIES = {
    '3d_thin': DataConfig('3d_thin', (128, 256, 256), (16, 16, 16), data_path=f'{DATA_ROOT}/3dthin'),
    '3d_mid': DataConfig('3d_mid', (64, 352, 352), (16, 16, 16), data_path=f'{DATA_ROOT}/3dmiddle'),
    '3d_thick': DataConfig('3d_thick', (24, 288, 288), (4, 16, 16), data_path=f'{DATA_ROOT}/3dthick'),
    '2d_grey': DataConfig('2d_grey', (512, 512), (16, 16), data_path=f'{DATA_ROOT}/2dgrey'),
    '2d_color': DataConfig('2d_color', (512, 512), (16, 16), num_channels=3, data_path=f'{DATA_ROOT}/2dcolor'),
    # '1d': DataConfig('1d', (4096, ), (16,), data_path='/data/datalake/1d'),
    # 'tab': DataConfig('tab', (2048, ), (1,), num_channels=2048, data_path='/data/datalake/0d')
}


class DataManager:
    def __init__(self, modality_list, total_dim=6144, seq_len=2048):
        self.n_modality = len(modality_list)
        self.modality_list = modality_list
        assert self.n_modality > 0
        self.modalities = [MODALITIES[m] for m in modality_list]

        self.total_dim = total_dim
        self.token_dims = [m.token_dim for m in self.modalities]

        ## enumerate different patch proj sizes
        for m_idx, m in enumerate(self.modalities):
            m.modality_idx = m_idx
            patch_proj_shape = (m.num_channels,) + tuple(m.patch_size)
            assert patch_proj_shape in PATCH_CONFIG
            m.patch_proj_idx = PATCH_CONFIG.index(patch_proj_shape)
            m.lo, m.hi = PATCH_RANGES[m.patch_proj_idx]

        self.seq_len = seq_len
        self.m_seq_lens = [m.seq_len for m in self.modalities]
        assert seq_len >= np.max(self.m_seq_lens)
        self.pad_lens = [seq_len - m.seq_len for m in self.modalities]

        # self.tab_idx = list(modality_list).index('tab') if 'tab' in modality_list else -1

    def __getitem__(self, idx):
        # data_config, start, end
        return self.modalities[idx]

    def get_modality_by_name(self, name):
        assert name in self.modality_list
        return self.modalities[self.modality_list.index(name)]

    def __len__(self):
        return len(self.modalities)


if __name__ == '__main__':
    manager = DataManager(['3d_thick', '3d_thin'])
    for m in manager:
        print(m[0].name)
